import pandas as pd
import re

from utils.loader import merge_xlsx_files


class SentenceExtractor:
    def __init__(self, vocab_excel_path, sheet_name='Sheet1'):
        """
        初始化工具类，加载词汇表
        :param vocab_excel_path: 词汇表 Excel 文件路径
        :param sheet_name: 工作表名称，默认为 'Sheet1'
        """
        # self.vocab_df = pd.read_excel(vocab_excel_path, sheet_name=sheet_name)
        self.vocab_df = merge_xlsx_files(vocab_excel_path)
        self.filtered_vocab = self.vocab_df.iloc[:, 0].tolist()
        self.sentences = []

    def filter_vocab_by_rank(self, min_rank=7000, max_rank=10000):
        """
        根据 rank 筛选词汇
        :param min_rank: 最小 rank
        :param max_rank: 最大 rank
        """
        self.filtered_vocab = self.vocab_df[
            (self.vocab_df['rank'] >= min_rank) &
            (self.vocab_df['rank'] <= max_rank)
            ]['lemma'].str.lower().tolist()

    def extract_sentences(self, text):
        """
        从文本中提取包含目标词汇的句子
        :param text: 输入文本（英文）
        """
        if self.filtered_vocab is None:
            raise ValueError("请先调用 filter_vocab_by_rank() 方法筛选词汇")

        # 使用正则表达式分割句子（简单按句号、问号、感叹号分割）
        sentences = re.split(r'(?<=[.!?])\s+', text)

        extracted_sentences = []
        for sentence in sentences:
            words = re.findall(r'\b\w+\b', sentence.lower())
            for vocab in self.filtered_vocab:
                if vocab in words:
                    print(vocab)
                    extracted_sentences.append({
                        'vocab': vocab,
                        'sentence': sentence.strip()
                    })
                    break  # 一个句子可能包含多个目标词，但每个句子只记录一次

        self.sentences = extracted_sentences

    def save_to_excel(self, output_path):
        """
        将提取的句子保存为 Excel 文件
        :param output_path: 输出 Excel 文件路径
        """
        if not self.sentences:
            print("没有提取到句子，不生成文件")
            return

        df_output = pd.DataFrame(self.sentences)
        df_output.to_excel(output_path, index=False)
        print(f"例句已保存至：{output_path}")


# 使用示例
if __name__ == "__main__":
    extractor = SentenceExtractor("lemmas_60k.xlsx")
    extractor.filter_vocab_by_rank(min_rank=10, max_rank=10000)

    # 示例文本（可替换为实际文本）
    sample_text = """
    This is a sample text. It contains words like 'example' and 'sample'. 
    Another sentence with the word 'test'. And one more with 'example' again. hold it.

    """

    extractor.extract_sentences(sample_text)
    extractor.save_to_excel("extracted_sentences.xlsx")